from random import uniform
import time
import numpy as np
from Crawl import crawl_vulnerability
import global_var

class Vulnerability:

    def __init__(self):
        self.id = "vul"
        # self.get_pkg_vul_dict()
        # self.check_vul()

    def get_pkg_vul_dict(self):
        pkg_list = np.load("./pkg_list.npy").tolist()
        print(len(pkg_list))
        all_dict = global_var.load_ver_dict()
        vul_dict = {}
        count = 1
        for index in range(len(pkg_list)):
            pkg = pkg_list[index]
            print("%s:%s" % (count, pkg))
            ver_dict = all_dict.get(pkg)
            if ver_dict == None:    # 没有版本信息的包
                vuls = crawl_vulnerability(pkg, "")
                if len(vuls) != 0:
                    vul_dict[pkg] = vuls
                sleep_time = uniform(0.1, 0.2)
                time.sleep(sleep_time)
            else:   # 有版本信息的包
                print("start get %s id from %s" % (len(ver_dict), pkg))
                for ver in ver_dict:
                    id = pkg + "@" + ver
                    vuls = crawl_vulnerability(pkg, ver)
                    if len(vuls) != 0:
                        vul_dict[id] = vuls
                    sleep_time = uniform(0.1, 0.2)
                    time.sleep(sleep_time)
                    print("%s ok!" % id)

            if count % 100 == 0 or count == len(pkg_list):
                file_count = count / 100 if count % 100 == 0 else count / 100 + 1
                print("already got %s pkg's vuls" % count)
                print("this bunch get %s valid" % len(vul_dict))
                np.save("./vuls/vuls_%s" % int(file_count), vul_dict)
                vul_dict.clear()

            count += 1
        print("all done")

    def get_proj_vul_dict(self):
        pkg_list = np.load("./pkg_list.npy").tolist()
        vul_dict = {}
        count = 1
        print("start")
        for pkg in pkg_list:
            print("%s:%s" % (count, pkg))
            vul_dict[pkg] = crawl_vulnerability(pkg, "")
            sleep_time = uniform(0.2, 0.4)
            time.sleep(sleep_time)
            if count % 100 == 0:
                print("already got %s pkg's vuls" % count)
            count += 1
        print("done")
        np.save("./vul_dict.npy", vul_dict)

    def check_vul(self):
        vul_dict = np.load("./vul_dict.npy", allow_pickle=True).item()
        count = 1
        for key in vul_dict:
            vuls = vul_dict[key]
            if len(vuls) != 0:
                print("%s:%s" % (count, key))
                for vul in vuls:
                    print(vul)
                count += 1
        return count

v = Vulnerability()